# cython: language_level=3
# -*- coding: utf-8 -*-
"""
轨迹处理和管理模块

该模块定义了强化学习中轨迹数据的处理、管理和训练流程。
包含轨迹类、轨迹池管理器和轨迹管理器等多个组件，支持完整的强化学习训练循环。

主要功能：
1. 轨迹数据的收集、存储和后处理
2. 奖励前推和GAE优势计算
3. 轨迹池的批量管理和训练触发
4. 多环境并行轨迹收集和管理
"""

import torch
from config import GlobalConfig
import numpy as np
from Common.trajectory_base import TRAJ_BASE
import copy
from Utils.tensor_ops import my_view


class trajectory(TRAJ_BASE):
    """
    轨迹类，继承自TRAJ_BASE

    扩展了基础轨迹类，添加了强化学习特定的功能，包括奖励处理、
    终止状态处理、优势计算等。

    Attributes:
        done (np.ndarray): 完成标志数组
        tobs: 终止观测状态
        team (int): 队伍标识（0：红队，1：蓝队）
        reference_track_name (str): 参考轨迹名称，默认为'reward'
    """

    def __init__(self, traj_limit, team, env_id):
        """
        初始化轨迹对象

        Args:
            traj_limit (int): 轨迹长度限制
            team (int): 队伍标识（0：红队，1：蓝队）
            env_id: 环境ID
        """
        super().__init__(traj_limit, env_id)
        self.done = None  # 完成标志数组
        self.tobs = None  # 终止观测状态
        self.team = team  # 队伍标识
        self.reference_track_name = 'reward'  # 参考轨迹名称（用于数据清理）

    def early_finalize(self):
        """
        提前结束轨迹

        当轨迹需要提前结束时调用，设置需要奖励回补标志。
        通常用于异常情况下的轨迹终止。
        """
        assert not self.readonly_lock, "轨迹已锁定，无法提前结束"
        self.need_reward_bootstrap = True

    def set_terminal_obs(self, tobs):
        """
        设置终止观测状态

        保存轨迹结束时的最终观测状态，用于价值函数的bootstrap计算。

        Args:
            tobs: 终止观测状态
        """
        self.tobs = copy.deepcopy(tobs)

    def set_done_state(self, shape):
        """
        设置完成状态数组

        创建指定形状的完成标志数组，全部初始化为True。

        Args:
            shape: 数组形状
        """
        self.done = torch.ones(*shape, dtype=torch.bool).unsqueeze(-1)

    def cut_tail(self):
        """
        截断轨迹尾部并清理无效数据

        扩展了基类的截断功能，增加了数据清理步骤：
        1. 调用基类方法截断预分配的尾部空间
        2. 根据参考轨迹（reward）检测和删除无效时间点
        3. 处理奖励前推，将无效时间点的奖励转移到前一个有效时间点
        4. 移除所有包含NaN的数据点

        这是轨迹数据预处理的关键步骤，确保训练数据的质量和一致性。
        """
        # 调用基类方法删除多余的预留空间
        super().cut_tail()
        TJ = lambda key: getattr(self, key)

        # 根据参考轨迹检测无效时间点
        reference_track = getattr(self, self.reference_track_name)

        # 处理需要奖励回补的情况（目前未启用）
        if self.need_reward_bootstrap:
            assert False, "此路径目前未启用，如需使用请检查实现"
            # 找到最后一个非NaN的位置
            T = np.where(~np.isnan(reference_track.squeeze()))[0][-1]
            self.boot_strap_value = {
                'bootstrap_value': TJ('value').squeeze()[T].copy(),
            }
            assert not hasattr(self, 'tobs')
            self.set_terminal_obs(TJ('g_obs')[T].copy())
            reference_track[T] = np.nan

        # 检测无效时间点：如果参考轨迹的任何维度为NaN，则该时间点无效
        p_invalid = np.isnan(my_view(reference_track, [0, -1])).any(axis=-1)
        p_valid = ~p_invalid

        # 如果所有时间点都无效，标记轨迹为废弃
        if p_invalid.all():
            self.deprecated_flag = True
            return

        # 奖励前推：将无效时间点的奖励转移到前一个有效时间点
        reward = TJ('reward')
        for i in reversed(range(self.time_pointer)):
            if p_invalid[i] and i != 0:  # 无效且不是第一个时间点
                reward[i - 1] += reward[i]  # 将奖励转移到前一个时间点
                reward[i] = np.nan  # 当前时间点奖励设为NaN
        setattr(self, 'reward', reward)

        # 移除所有无效时间点的数据
        for key in self.key_dict:
            setattr(self, key, TJ(key)[p_valid])

        return

    def reward_push_forward(self, dead_mask):
        """
        奖励前推处理

        根据死亡掩码将死亡时间点的奖励按折扣因子前推到之前的时间点。
        这是强化学习中处理终端状态奖励的常用技术。

        Args:
            dead_mask (np.ndarray): 死亡掩码，True表示该时间点对应的环境已死亡/结束

        工作原理：
        1. 从后向前遍历时间步
        2. 如果当前时间点死亡，将奖励按折扣因子gamma前推
        3. 将死亡时间点的奖励设为0
        """
        # 根据队伍获取对应的折扣因子
        if self.team == 0:
            gamma = GlobalConfig.AlgorithmConfig_Red.reward_forwarding_gamma
        else:
            gamma = GlobalConfig.AlgorithmConfig_Blue.reward_forwarding_gamma

        # 红队且启用奖励前推
        if self.team == 0 and GlobalConfig.AlgorithmConfig_Red.reward_forwarding:
            for i in reversed(range(self.time_pointer)):
                if i == 0: continue
                # 死亡时间点的奖励按gamma折扣前推，当前时间点奖励归零
                self.reward[i - 1] += np.where(dead_mask[i], self.reward[i] * gamma, 0)
                self.reward[i] = np.where(dead_mask[i], 0, self.reward[i])

        # 红队但未启用奖励前推（直接转移，无折扣）
        else:
            for i in reversed(range(self.time_pointer)):
                if i == 0: continue
                # 死亡时间点的奖励直接前推，无折扣
                self.reward[i - 1] += np.where(dead_mask[i], self.reward[i], 0)
                self.reward[i] = np.where(dead_mask[i], 0, self.reward[i])

        # 蓝队且启用奖励前推
        if self.team == 1 and GlobalConfig.AlgorithmConfig_Blue.reward_forwarding:
            for i in reversed(range(self.time_pointer)):
                if i == 0: continue
                self.reward[i - 1] += np.where(dead_mask[i], self.reward[i] * gamma, 0)
                self.reward[i] = np.where(dead_mask[i], 0, self.reward[i])

        # 蓝队但未启用奖励前推
        else:
            for i in reversed(range(self.time_pointer)):
                if i == 0: continue
                self.reward[i - 1] += np.where(dead_mask[i], self.reward[i], 0)
                self.reward[i] = np.where(dead_mask[i], 0, self.reward[i])

        return

    def finalize(self):
        """
        轨迹最终化处理

        对轨迹进行完整的最终化处理，包括数据验证、奖励处理、威胁计算等。
        这是轨迹完成收集后的标准处理流程。

        处理步骤：
        1. 验证数据完整性
        2. 检测死亡掩码
        3. 执行奖励前推
        4. 计算威胁度指标
        5. 准备GAE计算（已注释）
        """
        self.readonly_lock = True
        assert not self.deprecated_flag, "轨迹已废弃，无法最终化"
        TJ = lambda key: getattr(self, key)

        # 验证奖励数据完整性
        assert not np.isnan(TJ('reward')).any(), "奖励数据包含NaN值"

        # 检测死亡掩码：通过检查状态数据是否为NaN来确定死亡时间点
        tmp = np.isnan(my_view(self.state, [0, 0, -1]))  # 检查状态数据的最后一个维度
        dead_mask = tmp.all(-1)  # 如果所有子维度都是NaN，则认为该时间点死亡

        # 执行奖励前推处理
        self.reward_push_forward(dead_mask)

        # 计算威胁度指标：统计每个时间点之前活跃的智能体数量
        threat = np.zeros(shape=dead_mask.shape) - 1
        assert dead_mask.shape[0] == self.time_pointer

        for i in reversed(range(self.time_pointer)):
            if i + 1 < self.time_pointer:
                # 如果当前时间点或下一个时间点没有死亡，增加威胁度
                threat[:(i + 1)] += (~(dead_mask[i + 1] & dead_mask[i])).astype(int)
            elif i + 1 == self.time_pointer:
                # 最后一个时间点，如果未死亡则增加威胁度
                threat[:] += (~dead_mask[i]).astype(int)

        # 限制威胁度范围，防止异常值
        SAFE_LIMIT = 11
        threat = np.clip(threat, -1, SAFE_LIMIT)
        setattr(self, 'threat', np.expand_dims(threat, -1))

        # 注释：使用GAE计算回报（已修改，不再调用）
        # self.gae_finalize_return(reward_key='reward', value_key='value', new_return_name='return')
        return

    def gae_finalize_return(self, reward_key, value_key, new_return_name):
        """
        使用GAE（广义优势估计）计算回报

        实现GAE算法，用于计算强化学习中的优势函数和回报值。
        GAE结合了TD(0)和蒙特卡洛方法的优点，通过参数tau控制偏差-方差权衡。

        Args:
            reward_key (str): 奖励数据的键名
            value_key (str): 价值函数数据的键名
            new_return_name (str): 新计算的回报数据的键名

        算法步骤：
        1. 从后向前计算每个时间步的TD误差
        2. 使用指数衰减累积计算GAE
        3. 计算最终的回报值（优势 + 价值）
        """
        # 获取GAE参数
        if self.team == 0:
            gamma = GlobalConfig.AlgorithmConfig_Red.gamma  # 折扣因子
            tau = GlobalConfig.AlgorithmConfig_Red.tau     # GAE参数
        else:
            gamma = GlobalConfig.AlgorithmConfig_Blue.gamma
            tau = GlobalConfig.AlgorithmConfig_Blue.tau

        # 获取奖励和价值数据
        rewards = getattr(self, reward_key)
        value = getattr(self, value_key)
        length = rewards.shape[0]
        assert rewards.shape[0] == value.shape[0], "奖励和价值数据长度不匹配"

        # 确保奖励和价值数据维度一致
        if rewards.ndim == value.ndim - 1:
            rewards = np.expand_dims(rewards, -1)

        # 初始化回报数据数组
        setattr(self, new_return_name, np.zeros_like(value))
        self.key_dict.append(new_return_name)
        returns = getattr(self, new_return_name)

        # 设置bootstrap值（用于终端状态的价值估计）
        boot_strap = 0 if not self.need_reward_bootstrap else self.boot_strap_value['bootstrap_' + value_key]

        # 从后向前计算GAE和回报
        for step in reversed(range(length)):
            if step == (length - 1):  # 最后一帧
                value_preds_delta = rewards[step] + gamma * boot_strap - value[step]
                gae = value_preds_delta
            else:  # 其他帧
                value_preds_delta = rewards[step] + gamma * value[step + 1] - value[step]
                gae = value_preds_delta + gamma * tau * gae
            returns[step] = gae + value[step]


class TrajPoolManager(object):
    """
    轨迹池管理器

    负责管理轨迹池的生命周期，包括轨迹的最终化处理和批量训练准备。
    提供轨迹收集和训练之间的桥梁功能。
    """

    def __init__(self):
        """初始化轨迹池管理器"""
        self.cnt = 0  # 计数器

    def absorb_finalize_pool(self, pool):
        """
        吸收并最终化轨迹池

        对轨迹池中的所有轨迹进行最终化处理，包括截断、验证和清理。
        准备用于训练的轨迹数据。

        Args:
            pool (list): 轨迹池，包含多个轨迹对象

        Returns:
            tuple: (任务列表, 处理后的轨迹池)
        """
        # 对所有轨迹执行截断处理
        for traj_handle in pool:
            traj_handle.cut_tail()

        # 过滤掉废弃的轨迹
        pool = list(filter(lambda traj: not traj.deprecated_flag, pool))

        # 对剩余轨迹执行最终化处理
        for traj_handle in pool:
            traj_handle.finalize()

        self.cnt += 1
        task = ['train']  # 定义训练任务
        return task, pool


class TrajManagerBase(object):
    """
    轨迹管理器基类

    提供多环境并行轨迹收集的基础功能，包括轨迹创建、更新、
    数据验证和批量处理等核心功能。

    Attributes:
        n_env (int): 环境数量
        traj_limit (int): 单个轨迹的最大长度
        update_cnt (int): 更新计数器
        traj_pool (list): 轨迹池，存储已完成的轨迹
        registered_keys (list): 已注册的数据键列表
        live_trajs (list): 当前活跃的轨迹列表
        live_traj_frame (list): 每个活跃轨迹的当前帧数
        team (int): 队伍标识
    """

    def __init__(self, n_env, traj_limit, team):
        """
        初始化轨迹管理器基类

        Args:
            n_env (int): 并行环境数量
            traj_limit (int): 单个轨迹的最大长度限制
            team (int): 队伍标识（0：红队，1：蓝队）
        """
        self.n_env = n_env
        self.traj_limit = traj_limit
        self.update_cnt = 0
        self.traj_pool = []
        self.registered_keys = []  # 注册的轨迹数据键
        # 为每个环境创建活跃轨迹
        self.live_trajs = [trajectory(self.traj_limit, team, env_id=i) for i in range(self.n_env)]
        self.live_traj_frame = [0 for _ in range(self.n_env)]
        self._traj_lock_buf = None  # 轨迹锁定缓冲区
        self.patience = 1000  # 完整性检查耐心值
        self.team = team  # 队伍标识

    def __check_integraty(self, traj_frag):
        """
        检查轨迹片段的数据完整性

        验证轨迹片段中包含的数据键是否一致，注册新的数据键。
        用于确保数据收集过程中的一致性。

        Args:
            traj_frag (dict): 轨迹片段数据字典
        """
        if self.patience < 0:
            return  # 停止检查以节省时间

        self.patience -= 1

        # 注册新的数据键
        for key in traj_frag:
            if key not in self.registered_keys and (not key.startswith('_')):
                self.registered_keys.append(key)

        # 验证所有已注册键的存在性
        for key in self.registered_keys:
            assert key in traj_frag, f"轨迹片段中缺少键: {key}"

    def __check_key_indict(self, dictionary, key):
        """
        检查字典中指定键的存在性并返回形状

        Args:
            dictionary (dict): 数据字典
            key (str): 要检查的键名

        Returns:
            tuple: 数据形状

        Raises:
            AssertionError: 如果键不存在
        """
        assert key in dictionary, f"轨迹中未找到键 '{key}'"
        return dictionary[key].shape

    def batch_update(self, traj_frag):
        """
        批量更新轨迹数据

        处理来自多个环境的轨迹片段数据，包括数据分发、轨迹管理、
        完成检测和新轨迹创建等功能。

        Args:
            traj_frag (dict): 轨迹片段数据，包含来自多个环境的数据
        """
        self.__check_integraty(traj_frag)

        # 提取特殊标志
        done = traj_frag.pop('_DONE_')  # 完成标志
        skip = traj_frag.pop('_SKIP_')  # 跳过/冻结标志
        tobs = traj_frag.pop('_TOBS_')  # 终止观测

        # 转换单个布尔值为列表
        if isinstance(done, bool):
            done = [done for i in range(self.n_env)]
        if isinstance(skip, bool):
            skip = [skip for i in range(self.n_env)]

        n_active = sum(~skip)  # 活跃环境数量

        # 分发数据到各个环境
        cnt = 0
        for env_i in range(self.n_env):
            if skip[env_i]:
                continue  # 跳过冻结的环境

            frag_index = cnt
            cnt += 1
            env_index = env_i
            traj_handle = self.live_trajs[env_index]

            # 添加完成标志数组（2024/08修改）
            action_shape = self.__check_key_indict(traj_frag, "action")
            if done.all():
                done_state = np.ones(action_shape, dtype=np.bool8)
            else:
                done_state = np.zeros(action_shape, dtype=np.bool8)
            traj_frag.update({'done': done_state})

            # 将数据存储到轨迹中
            for key in traj_frag:
                self.traj_remember(traj_handle, key=key, content=traj_frag[key],
                                 frag_index=frag_index, n_active=n_active)

            # 更新轨迹状态
            self.live_traj_frame[env_index] += 1
            traj_handle.time_shift()

            # 处理轨迹完成
            if done[env_i]:
                assert tobs[env_i] is not None, "完成轨迹必须有终止观测"
                traj_handle.set_terminal_obs(tobs[env_i])
                self.traj_pool.append(traj_handle)  # 添加到轨迹池
                # 创建新的活跃轨迹
                self.live_trajs[env_index] = trajectory(self.traj_limit, env_id=env_index, team=self.team)
                self.live_traj_frame[env_index] = 0

    def traj_remember(self, traj, key, content, frag_index, n_active):
        """
        将数据存储到轨迹中

        支持嵌套数据结构（字典）的递归存储。

        Args:
            traj: 轨迹对象
            key (str): 数据键名
            content: 数据内容
            frag_index (int): 片段索引
            n_active (int): 活跃环境数量
        """
        if content is None:
            traj.remember(key, None)
        elif isinstance(content, dict):
            # 递归处理嵌套字典
            for sub_key in content:
                self.traj_remember(traj, "".join((key, ">", sub_key)),
                                 content=content[sub_key], frag_index=frag_index,
                                 n_active=n_active)
        else:
            assert n_active == len(content), f"数据长度错误: {n_active} != {len(content)}"
            traj.remember(key, content[frag_index])


class TrajectoryManager(TrajManagerBase):
    """
    轨迹管理器，继承自TrajManagerBase

    提供完整的轨迹管理功能，包括训练触发、批量处理、
    梯度更新等高级功能。

    Attributes:
        trainer_hook (callable): 训练钩子函数
        traj_limit (int): 轨迹长度限制
        num_training_trajectories (int): 训练所需的最小轨迹数量
        pool_manager (TrajPoolManager): 轨迹池管理器
    """

    def __init__(self, num_env, max_trajectory_data, train_model_func, team):
        """
        初始化轨迹管理器

        Args:
            num_env (int): 环境数量
            max_trajectory_data (int): 最大轨迹数据长度
            train_model_func (callable): 训练模型函数
            team (int): 队伍标识
        """
        super().__init__(num_env, max_trajectory_data, team)
        self.trainer_hook = train_model_func  # 训练钩子函数
        self.traj_limit = max_trajectory_data

        # 根据队伍设置训练轨迹数量要求
        if self.team == 0:
            self.num_training_trajectories = GlobalConfig.AlgorithmConfig_Red.num_training_trajectories
        else:
            self.num_training_trajectories = GlobalConfig.AlgorithmConfig_Blue.num_training_trajectories

        self.pool_manager = TrajPoolManager()

    def update(self, traj_frag, index):
        """
        更新指定环境的轨迹数据

        Args:
            traj_frag (dict): 轨迹片段数据
            index (list): 要更新的环境索引列表
        """
        assert traj_frag is not None

        # 对每个指定环境更新轨迹数据
        for j, env_i in enumerate(index):
            traj_handle = self.live_trajs[env_i]

            for key in traj_frag:
                assert traj_frag[key] is not None, f"数据为空: {key}"

                if isinstance(traj_frag[key], dict):
                    # 处理嵌套字典数据
                    for sub_key in traj_frag[key]:
                        content = traj_frag[key][sub_key][j]
                        traj_handle.remember(key + ">" + sub_key, content)
                else:
                    # 处理普通数据
                    content = traj_frag[key][j]
                    traj_handle.remember(key, content)

            self.live_traj_frame[env_i] += 1
            traj_handle.time_shift()

        return

    def feed_traj_framedata(self, traj_frag, require_hook=False):
        """
        喂入轨迹帧数据（函数入口）

        处理轨迹数据的喂入，支持锁定机制以处理不完整的轨迹片段。

        Args:
            traj_frag (dict): 轨迹片段数据
            require_hook (bool): 是否需要钩子锁定，默认False

        Returns:
            callable or None: 如果需要钩子则返回解锁函数
        """
        # 确保没有锁定缓冲区
        assert self._traj_lock_buf is None

        if require_hook:
            # 不完整的轨迹片段，需要锁定等待更多数据
            assert '_SKIP_' in traj_frag
            assert '_DONE_' not in traj_frag
            assert 'reward' not in traj_frag
            self._traj_lock_buf = traj_frag
            return self.unlock_fn
        else:
            # 完整的轨迹片段，直接处理
            assert '_DONE_' in traj_frag
            assert '_SKIP_' in traj_frag
            self.batch_update(traj_frag=traj_frag)
            return

    def clear_traj_pool(self):
        """
        清空轨迹池

        处理轨迹池中的所有轨迹并清空，用于训练前的数据准备。
        """
        print(f'执行更新 {self.update_cnt}')
        _, self.traj_pool = self.pool_manager.absorb_finalize_pool(pool=self.traj_pool)
        self.traj_pool = []
        return self.update_cnt

    def train_with_trajectory_data(self):
        """
        使用轨迹数据进行训练

        最终化轨迹池数据并执行模型训练。

        Returns:
            int: 更新计数器
        """
        print(f'执行模型训练: {self.update_cnt + 1}')

        # 最终化轨迹池数据
        current_task_l, self.traj_pool = self.pool_manager.absorb_finalize_pool(pool=self.traj_pool)

        # 对每个任务执行训练
        for current_task in current_task_l:
            ppo_update_cnt = self.trainer_hook(self.traj_pool, current_task)

        # 清空轨迹池并更新计数器
        self.traj_pool = []
        self.update_cnt += 1
        return self.update_cnt

    def is_ready_to_train(self):
        """
        检查是否准备好进行训练

        Returns:
            bool: 如果轨迹池中有足够的轨迹数据则返回True
        """
        if len(self.traj_pool) >= self.num_training_trajectories:
            print(f'{len(self.traj_pool)} 个轨迹数据，轨迹限制: {self.traj_limit}')
            return True
        else:
            return False

    def unlock_fn(self, traj_frag):
        """
        解锁函数，用于处理锁定的轨迹片段

        将锁定的轨迹片段与新数据合并并处理。

        Args:
            traj_frag (dict): 新的轨迹片段数据
        """
        assert self._traj_lock_buf is not None
        traj_frag.update(self._traj_lock_buf)
        self._traj_lock_buf = None
        assert '_DONE_' in traj_frag
        assert '_SKIP_' in traj_frag
        self.batch_update(traj_frag=traj_frag)